#include <iostream>
#include <cuda_runtime.h>
#include <chrono>

using namespace std::chrono;

const int DATASIZE = 5000000;
const int BLOCK_SIZE = 256; 
const int UINT32_BITS = 32;  // 32位无符号整数的位数
// const int DATASIZE = 11;
// const int BLOCK_SIZE = 3; 
// const int UINT32_BITS = 1;


bool validate(unsigned int *a, unsigned int* b, int n){
    for (int i = 0; i < n; i++) {
        if(a[i] != b[i]){
            return false;
        }
    }
    return true;
}

void print_array(unsigned int* arr, int n, const char* label) {
    if(n > 20) return;
    printf("\n%s\n", label);
    for (int i = 0; i < n; i++) {
        printf("%d  ", arr[i]);
    }
    printf("\n");
}

void cpu_radix_sort_with_index(unsigned int* arr, unsigned int* arr_pos,
                               unsigned int* sorted_arr, unsigned int* original_indices,
                               int n) {
    unsigned int* values_curr = new unsigned int[n];
    unsigned int* values_next = new unsigned int[n];
    unsigned int* pos_curr = new unsigned int[n];
    unsigned int* pos_next = new unsigned int[n];
    std::copy(arr, arr + n, values_curr);
    std::copy(arr_pos, arr_pos + n, pos_curr);
    for (int bit = 0; bit < UINT32_BITS; bit++) {
        int count[2] = {0, 0};
        for (int i = 0; i < n; i++) {
            int digit = (values_curr[i] >> bit) & 1;
            count[digit]++;
        }
        int offset[2] = {0, count[0]};
        for (int i = 0; i < n; i++) {
            int digit = (values_curr[i] >> bit) & 1;
            int dest_idx = offset[digit]++;
            values_next[dest_idx] = values_curr[i];
            pos_next[dest_idx] = pos_curr[i];
        }
        std::swap(values_curr, values_next);
        std::swap(pos_curr, pos_next);
    }
    // 循环结束后，最终结果在 values_curr 和 pos_curr 中
    // 因为每次 swap 后，curr 总是指向最新结果
    std::copy(values_curr, values_curr + n, sorted_arr);
    std::copy(pos_curr, pos_curr + n, original_indices);

    delete[] values_curr;
    delete[] values_next;
    delete[] pos_curr;
    delete[] pos_next;
}

// 扫描类型枚举
enum ScanType {
 INCLUSIVE,  // 包含式扫描(包含当前元素)
 EXCLUSIVE   // 排除式扫描(不包含当前元素)
};

__global__
void naive_histogram(unsigned int* const d_histo, int numBins,
                     unsigned int* const d_in, int size, int shift) {
  int i = threadIdx.x + blockDim.x * blockIdx.x;
  if (i >= size) return;

  unsigned int bit = ((d_in[i] >> shift) & 1u);
  atomicAdd(&d_histo[bit], 1);
}


// Hillis-Steele并行扫描核函数
__global__
void scan_hillis_steele(unsigned int* d_out, const unsigned int* d_in, int size, ScanType scan_type) {
  extern __shared__ unsigned int shared_data[];  // 两倍 blocksize 大小,双缓冲避免读写冲突
  int tid = threadIdx.x;  // 线程ID
  int globalId = blockDim.x * blockIdx.x + threadIdx.x;  // 全局ID
  if (globalId >= size) return;  // 边界检查

  // 根据扫描类型初始化
  if (scan_type == INCLUSIVE) {
    shared_data[tid] = d_in[globalId];  // 包含式
  } else {
    shared_data[tid] = (globalId == 0 ? 0 : d_in[globalId-1]);  // 排除式
  }
  __syncthreads();  // 同步线程

  // Hillis-Steele扫描算法
  int pin = 0, pout = 1;  // 双缓冲索引
  for (int s = 1; s < blockDim.x; s <<= 1) {  // 步长倍增
    if (tid < s) {
      shared_data[pout * blockDim.x + tid] = shared_data[pin * blockDim.x + tid];
    }
    if (tid + s < blockDim.x) {
      shared_data[pout * blockDim.x + tid + s] = 
        shared_data[pin * blockDim.x + tid] + shared_data[pin * blockDim.x + tid + s];
    }
    __syncthreads();
    pin = 1 - pin; pout = 1 - pout;  // 交换缓冲区
  }
  d_out[globalId] = shared_data[pin * blockDim.x + tid];  // 写入结果
}

// 块扫描相加核函数
__global__
void block_scan_add_block_sum(unsigned int* d_out, const unsigned int* const block_sum, 
                             const unsigned int* const d_scan, int size) {
  int i = blockIdx.x * blockDim.x + threadIdx.x;  // 全局索引
  if (i >= size) return;  // 边界检查
  d_out[i] = block_sum[blockIdx.x] + d_scan[i];  // 合并结果
}

// 零位映射核函数
__global__
void zero_digit_mapper(unsigned int* d_out, unsigned int* const d_inputVals,
                      size_t numElems, int shift) {
  int i = blockIdx.x * blockDim.x + threadIdx.x;  // 全局索引
  if (i >= numElems) return;  // 边界检查
  d_out[i] = !((d_inputVals[i] >> shift) & 1u);  // 检查是否为0位,为0则标记为1
}

// 初始化块和核函数
__global__ 
void init_block_sum(unsigned int* block_sum, const unsigned int* const block_level_scan, int block_num) {
  int i = blockIdx.x * blockDim.x + threadIdx.x;  // 全局索引
  if (i < block_num) {  // 边界检查
    block_sum[i] = (i == 0) ? 0 : block_level_scan[i * blockDim.x - 1];  // 计算块前缀和
  }
}

// 递归扫描函数(如果递归次数过多，存在内存问题)
void scan(unsigned int* d_out, const unsigned int* d_in, int size, ScanType scan_type, int recursive){
  if (size <= 0) return;  // 边界检查
  // printf("recursive = %d\n", recursive);

  // 1. 块级扫描
  int block_num = (size + BLOCK_SIZE - 1) / BLOCK_SIZE;  // 计算块数
  scan_hillis_steele<<<block_num, BLOCK_SIZE, 2 * BLOCK_SIZE * sizeof(unsigned int)>>>(d_out, d_in, size, scan_type);
  cudaDeviceSynchronize();

  // unsigned int* h_out = new unsigned int[size];
  // cudaMemcpy(h_out, d_out, size * sizeof(unsigned int), cudaMemcpyDeviceToHost);
  // print_array(h_out, size, "1-scan:");
  // delete[] h_out;
  
  if (block_num == 1) {
    return;  // 单块直接返回
  }

  // 2. 计算块和并递归扫描
  unsigned int* d_block_sum;
  cudaMalloc(&d_block_sum, sizeof(unsigned int) * block_num);
  int block_num2 = (block_num + BLOCK_SIZE-1)/BLOCK_SIZE;
  init_block_sum<<<block_num2, BLOCK_SIZE>>>(d_block_sum, d_out, block_num);

  // unsigned int* h_block_sum = new unsigned int[block_num];
  // cudaMemcpy(h_block_sum, d_block_sum, block_num * sizeof(unsigned int), cudaMemcpyDeviceToHost);
  // print_array(h_block_sum, block_num, "2-init:");
  // delete[] h_block_sum;

  scan(d_block_sum, d_block_sum, block_num, INCLUSIVE, recursive + 1);

  // unsigned int* h_block_sum2 = new unsigned int[block_num];
  // cudaMemcpy(h_block_sum2, d_block_sum, block_num * sizeof(unsigned int), cudaMemcpyDeviceToHost);
  // print_array(h_block_sum2, block_num, "2-scan:");
  // delete[] h_block_sum2;

  // 3. 合并扫描结果
  block_scan_add_block_sum<<<block_num, BLOCK_SIZE>>>(d_out, d_block_sum, d_out, size);
  cudaDeviceSynchronize(); 

  // unsigned int* h_val = new unsigned int[size];
  // cudaMemcpy(h_val, d_out, size * sizeof(unsigned int), cudaMemcpyDeviceToHost);
  // print_array(h_val, size, "3-addblock:");
  // delete[] h_val;

  cudaFree(d_block_sum);  // 释放内存
}

// 数据移动核函数
__global__ 
void move_data(unsigned int* d_out_vals, unsigned int* d_out_pos,
              unsigned int* const d_in_values, unsigned int* const d_in_pos,
              unsigned int* d_zero_pos, unsigned int* d_histo,
              int shift, int size) {
  int i = threadIdx.x + blockDim.x * blockIdx.x;  // 全局索引
  if (i >= size) return;  // 边界检查

  unsigned int tmp_value = d_in_values[i];  // 临时存储值
  unsigned int tmp_pos = d_in_pos[i];  // 临时存储位置
  __syncthreads();  // 同步线程

  unsigned int num_zeros = d_histo[0];  // 获取0位总数量

  // 根据当前位值移动数据
  if(((d_in_values[i] >> shift) & 0x1u) == 0) {
    d_out_vals[d_zero_pos[i]] = tmp_value;  // 0位放置在前部
    d_out_pos[d_zero_pos[i]] = tmp_pos;
  } else {
    // 1位放置在所有0位的后部
    d_out_vals[i - d_zero_pos[i] + num_zeros] = tmp_value;
    d_out_pos[i - d_zero_pos[i] + num_zeros] = tmp_pos;
  }
}

void gpu_radix_sort(unsigned int* const d_inputVals,
               unsigned int* const d_inputPos,
               unsigned int* const d_outputVals,
               unsigned int* const d_outputPos,
               const size_t numElems){
  // 分配设备内存用于存储0位标记
  unsigned int* d_zero;
  cudaMalloc(&d_zero, sizeof(unsigned int)*numElems);

  // 分配设备内存用于直方图(2 bins)
  unsigned int* d_histo;
  cudaMalloc(&d_histo, sizeof(unsigned int)*2);

  // 初始化源指针
  unsigned int* d_vals_src = d_inputVals;
  unsigned int* d_pos_src = d_inputPos;

  // 初始化目标指针
  unsigned int* d_vals_dst_tmp = d_outputVals;
  unsigned int* d_pos_dst_tmp = d_outputPos;

  int gridsize = (numElems + BLOCK_SIZE - 1)/BLOCK_SIZE;

  // 对32位无符号整数的每一位进行处理(从最低位开始)
  for (int shift=0; shift < UINT32_BITS; shift++) {
    // 重置直方图
    cudaMemset(d_histo, 0, sizeof(unsigned int)*2);
    
    // 计算当前位的直方图(0和1的计数)
    naive_histogram<<<gridsize, BLOCK_SIZE>>>(d_histo, 2, d_vals_src, numElems, shift);
    cudaDeviceSynchronize();

    // unsigned int* h_histo = new unsigned int[2];
    // cudaMemcpy(h_histo, d_histo, 2 * sizeof(unsigned int), cudaMemcpyDeviceToHost);
    // print_array(h_histo, 2, "histo:");
    // delete[] h_histo;

    // 标记当前位为0的元素,为0则标记为1
    zero_digit_mapper<<<gridsize, BLOCK_SIZE>>>(d_zero, d_vals_src, numElems, shift);
    cudaDeviceSynchronize(); 

    // unsigned int* h_zero = new unsigned int[numElems];
    // cudaMemcpy(h_zero, d_zero,  numElems * sizeof(unsigned int), cudaMemcpyDeviceToHost);
    // print_array(h_zero, numElems, "zero:");
    

    // 对0位标记进行前缀和扫描(非包含)
    scan(d_zero, d_zero, numElems, EXCLUSIVE, 1);

    // cudaMemcpy(h_zero, d_zero,  numElems * sizeof(unsigned int), cudaMemcpyDeviceToHost);
    // print_array(h_zero, numElems, "zero prefixsum:");
    // delete[] h_zero;
    
    // 根据扫描结果移动数据到正确位置
    move_data<<<gridsize, BLOCK_SIZE>>>(d_vals_dst_tmp, d_pos_dst_tmp,
                                                                  d_vals_src, d_pos_src, d_zero, d_histo, shift, numElems);
    cudaDeviceSynchronize(); 

    // unsigned int* h_out = new unsigned int[numElems];
    // unsigned int* h_out_pos = new unsigned int[numElems];
    // cudaMemcpy(h_out, d_vals_dst_tmp,  numElems * sizeof(unsigned int), cudaMemcpyDeviceToHost);
    // cudaMemcpy(h_out_pos, d_pos_dst_tmp,  numElems * sizeof(unsigned int), cudaMemcpyDeviceToHost);
    // print_array(h_out, numElems, "move data:");
    // print_array(h_out_pos, numElems, "move pos:");
    // delete[] h_out;delete[] h_out_pos;

    // 缓冲区交换
    std::swap(d_vals_src, d_vals_dst_tmp);
    std::swap(d_pos_src, d_pos_dst_tmp);
  }

  // 因为每次 swap 后，src 总是指向最新结果
  if(UINT32_BITS % 2 == 0){
    cudaMemcpy(d_outputVals, d_vals_src,  numElems * sizeof(unsigned int), cudaMemcpyDeviceToDevice);
    cudaMemcpy(d_outputPos, d_pos_src,  numElems * sizeof(unsigned int), cudaMemcpyDeviceToDevice);
  }

  // 释放临时内存
  cudaFree(d_histo);
  cudaFree(d_zero);
}

int main() {
    // unsigned int* h_in = new unsigned int[DATASIZE] {5, 3, 7, 1, 9, 2 , 8, 4, 6, 0, 3};
    // unsigned int* h_in_pos = new unsigned int[DATASIZE] {0, 1, 2, 3, 4, 5 , 6, 7, 8, 9, 10};
    unsigned int* h_in = new unsigned int[DATASIZE];
    unsigned int* h_in_pos = new unsigned int[DATASIZE];
    unsigned int* h_out = new unsigned int[DATASIZE];
    unsigned int* h_out_pos = new unsigned int[DATASIZE];
    unsigned int* cpu_out = new unsigned int[DATASIZE];
    unsigned int* cpu_out_pos = new unsigned int[DATASIZE];

    srand(100);
    for (int i = 0; i < DATASIZE; ++i) {
        h_in[i] = rand() % (DATASIZE + DATASIZE / 10);
        h_in_pos[i] = i;
    } 

    auto start_cpu = high_resolution_clock::now();
    cpu_radix_sort_with_index(h_in, h_in_pos, cpu_out, cpu_out_pos, DATASIZE);
    auto end_cpu = high_resolution_clock::now();
    auto cpu_time = duration_cast<milliseconds>(end_cpu - start_cpu).count();
    
    auto start = high_resolution_clock::now();
    unsigned int *d_in, *d_in_pos, *d_out, *d_out_pos;
    cudaMalloc(&d_in, DATASIZE * sizeof(unsigned int));
    cudaMalloc(&d_in_pos, DATASIZE * sizeof(unsigned int));
    cudaMalloc(&d_out, DATASIZE * sizeof(unsigned int));
    cudaMalloc(&d_out_pos, DATASIZE * sizeof(unsigned int));

    cudaMemcpy(d_in, h_in, DATASIZE * sizeof(unsigned int), cudaMemcpyHostToDevice);
    cudaMemcpy(d_in_pos, h_in_pos, DATASIZE * sizeof(unsigned int), cudaMemcpyHostToDevice);
    cudaMemcpy(d_out, h_out, DATASIZE * sizeof(unsigned int), cudaMemcpyHostToDevice);
    cudaMemcpy(d_out_pos, h_out_pos, DATASIZE * sizeof(unsigned int), cudaMemcpyHostToDevice);

    gpu_radix_sort(d_in, d_in_pos, d_out, d_out_pos, DATASIZE);

    cudaMemcpy(h_out,d_out, DATASIZE * sizeof(unsigned int), cudaMemcpyDeviceToHost);
    cudaMemcpy(h_out_pos, d_out_pos, DATASIZE * sizeof(unsigned int), cudaMemcpyDeviceToHost);
    cudaFree(d_in);cudaFree(d_in_pos);cudaFree(d_out);cudaFree(d_out_pos);
    auto end = high_resolution_clock::now();
    auto gpu_time = duration_cast<milliseconds>(end - start).count();

    // print_array(h_in, DATASIZE, "originally value:");
    // print_array(cpu_out, DATASIZE, "cpu sorted value:");
    // print_array(h_out, DATASIZE, "gpu sorted value:");
    // print_array(cpu_out_pos, DATASIZE, "cpu sorted index:");
    // print_array(h_out_pos, DATASIZE, "gpu sorted index:");

    printf("Sort value: %s\n", validate(h_out, cpu_out, DATASIZE) ? "CORRECT" : "INCORRECT");
    printf("Sort index: %s\n", validate(h_out_pos, cpu_out_pos, DATASIZE) ? "CORRECT" : "INCORRECT");
    std::cout << "Data size: " << DATASIZE << " elements" << std::endl;
    std::cout << "GPU  time: " << gpu_time << " ms" << std::endl;
    std::cout << "CPU  time: " << cpu_time << " ms" << std::endl;
    
    delete[] h_in;
    delete[] h_in_pos;
    delete[] cpu_out;
    delete[] cpu_out_pos;
    return 0;
}

